package cn.kinoko.common.utils;

import java.util.Arrays;
import java.util.Collection;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.function.Consumer;

/**
 * 异步工具类
 *
 * @author kinoko
 * @date 2023/7/13 10:07
 */
public class CompletableFutureUtil {

    /**
     * 循环执行异步方法
     *
     * @param collection 执行的集合
     * @param function   对集合对象的操作
     * @param <T>        集合对象类型
     * @throws InterruptedException e
     */
    public static <T> void runAsyncForEach(Collection<T> collection, Consumer<T> function) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(collection.size());
        for (T obj : collection) {
            CompletableFuture.runAsync(() -> {
                try {
                    function.accept(obj);
                } catch (Exception e) {
                    throw new RuntimeException(e);
                } finally {
                    latch.countDown();
                }
            });
        }
        latch.await();
    }

    /**
     * 循环执行异步方法
     *
     * @param collection 执行的集合
     * @param function   对集合对象的操作
     * @param executor   线程池
     * @param <T>        集合对象类型
     * @throws InterruptedException e
     */
    public static <T> void runAsyncForEach(Collection<T> collection, Consumer<T> function, Executor executor) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(collection.size());
        for (T obj : collection) {
            CompletableFuture.runAsync(() -> {
                try {
                    function.accept(obj);
                } catch (Exception e) {
                    throw new RuntimeException(e);
                } finally {
                    latch.countDown();
                }
            }, executor);
        }
        latch.await();
    }

    /**
     * 循环执行异步方法并阻塞
     *
     * @param tasks 执行的集合
     * @throws ExecutionException   e
     * @throws InterruptedException e
     */
    public static void runAsyncAllAndWait(Runnable... tasks) throws ExecutionException, InterruptedException {
        if (tasks.length == 0) return;
        CompletableFuture<?>[] futures = Arrays.stream(tasks)
                .parallel().map(CompletableFuture::runAsync)
                .toArray(CompletableFuture[]::new);
        CompletableFuture.allOf(futures).get();
    }

    /**
     * 循环执行异步方法并阻塞
     *
     * @param executor 线程池
     * @param tasks 执行的集合
     * @throws ExecutionException   e
     * @throws InterruptedException e
     */
    public static void runAsyncAllAndWait(Executor executor, Runnable... tasks) throws ExecutionException, InterruptedException {
        if (tasks.length == 0) return;
        CompletableFuture<?>[] futures = Arrays.stream(tasks)
                .parallel().map(task -> CompletableFuture.runAsync(task, executor))
                .toArray(CompletableFuture[]::new);
        CompletableFuture.allOf(futures).get();
    }

}